""" train and test dataset

author baiyu
"""
import os
import sys
import pickle

from skimage import io
import matplotlib.pyplot as plt
import numpy
import torch
from torch.utils.data import Dataset
import random
from PIL import Image
import cv2
import numpy as np

def cv_imread(file_path = ""):
    img_mat=cv2.imdecode(np.fromfile(file_path,dtype=np.uint16),-1)
    # img_mat = img_mat.astype('float32') / 65535.0
    return img_mat
# D:\\dingding\\ImageTrain\\train\\mine\\1.284705_avg0_170_8_18_180_1.tif
class StoneTrain(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """

    def __init__(self, path, transform=None):
        self.data = []
        self.path = path
       # self.pathmine = 'G:\\dingding\\data\\ImageTrain\\trainPos'# os.path.join(path, 'mine')
       # self.pathwaste ='G:\\dingding\\data\\ImageTrain\\trainNeg' # os.path.join(path, 'waste')
        self.pathmine = os.path.join(path, 'mine')
        self.pathwaste = os.path.join(path, 'waste')
        # readline paper
        for file in os.listdir(self.pathmine):
            if '.tif' in file:
                result = (0, file)  # mine 0
                self.data.append(result)

        for file in os.listdir(self.pathwaste):
            if '.tif' in file:
                result = (1, file)    # waste 1
                self.data.append(result)

        random.shuffle(self.data)
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        # label = self.data['fine_labels'.encode()][index]
        # r = self.data['data'.encode()][index, :1024].reshape(32, 32)
        # g = self.data['data'.encode()][index, 1024:2048].reshape(32, 32)
        # b = self.data['data'.encode()][index, 2048:].reshape(32, 32)
        label = self.data[index][0]
        image = []
        if label == 1:  # 0
            path = os.path.join(self.path, self.pathwaste, self.data[index][1])
            # image = cv2.imread(path)
            image = cv_imread(path)
            image = image[..., ::-1]
           #  image = Image.open(path)# numpy.dstack((r, g, b))
        else:
            path = os.path.join(self.path, self.pathmine, self.data[index][1])

            # image = cv2.imread(path)
            image = cv_imread(path)
            '''
            if image is None:
                print('你是煞笔')
                print(path)
            '''
            image = image[..., ::-1]
            # image = Image.open(path)
       #  cv2.imshow("demo", image)
       #  cv2.waitKey(0)
        # image.show()
        image = Image.fromarray(np.uint8(image))
        image = image.resize((96, 96))
        image = image.copy()
        image = np.uint8(image)
        if self.transform:
            image = self.transform(image)

        # image = torch.tensor(image)
        return image, label


class StoneTest(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """

    def __init__(self, path, transform=None):
        self.data = []
        self.path = path
        self.pathmine = os.path.join(path, 'mine')
        self.pathwaste = os.path.join(path, 'waste')
        # readline paper
        for file in os.listdir(self.pathmine):
            if '.tif' in file:
                result = (0, file)
                self.data.append(result)

        for file in os.listdir(self.pathwaste):
            if '.tif' in file:
                result = (1, file)
                self.data.append(result)

        random.shuffle(self.data)
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        label = self.data[index][0]
        image = []
        if label == -1:
            path = os.path.join(self.path, self.pathwaste, self.data[index][1])
            # image = cv2.imread(path)
            image = cv_imread(path)
            image = image[..., ::-1]
        #  image = Image.open(path)# numpy.dstack((r, g, b))
        else:
            path = os.path.join(self.path, self.pathmine, self.data[index][1])
            # image = cv2.imread(path)
            image = cv_imread(path)
            image = image[..., ::-1]
            # image = Image.open(path)
        #  cv2.imshow("demo", image)
        #  cv2.waitKey(0)
        # image.show()
        image = Image.fromarray(np.uint8(image))
        image = image.resize((96, 96))
        if self.transform:
            image = self.transform(image)
        return  image, label

# Stone test divide two set in our public set
class StoneTestMine(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """
    def __init__(self, path, transform=None):
        self.data = []
        self.path = path
        self.pathmine = os.path.join(path, 'mine')
        # readline paper
        for file in os.listdir(self.pathmine):
            if '.tif' in file:
                result = (1, file)
                self.data.append(result)
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        label = self.data[index][0]
        image = []
        path = os.path.join(self.path, self.pathmine, self.data[index][1])
        # image = cv2.imread(path)
        image = cv_imread(path)
        image = image[..., ::-1].copy()
        image = np.uint8(image)
        if self.transform:
            image = self.transform(image)
        return  image, label


class StoneTestWaste(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """
    def __init__(self, path, transform=None):
        self.data = []
        self.path = path
        self.pathwaste = os.path.join(path, 'waste')

        for file in os.listdir(self.pathwaste):
            if '.tif' in file:
                result = (-1, file)
                self.data.append(result)
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        label = self.data[index][0]
        image = []
        path = os.path.join(self.path, self.pathwaste, self.data[index][1])
            # image = cv2.imread(path)
        image = cv_imread(path)
        image = image[..., ::-1].copy()
        image = np.uint8(image)

        if self.transform:
            image = self.transform(image)
        return  image, label


'''  before code 
class CIFAR100Train(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """

    def __init__(self, path, transform=None):
        #if transform is given, we transoform data using
        with open(os.path.join(path, 'train'), 'rb') as cifar100:
            self.data = pickle.load(cifar100, encoding='bytes')
        self.transform = transform

    def __len__(self):
        return len(self.data['fine_labels'.encode()])

    def __getitem__(self, index):
        label = self.data['fine_labels'.encode()][index]
        r = self.data['data'.encode()][index, :1024].reshape(32, 32)
        g = self.data['data'.encode()][index, 1024:2048].reshape(32, 32)
        b = self.data['data'.encode()][index, 2048:].reshape(32, 32)
        image = numpy.dstack((r, g, b))
        
        if self.transform:
            image = self.transform(image)
        return label, image


class CIFAR100Test(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """

    def __init__(self, path, transform=None):
        # readline paper
        with open(os.path.join(path, 'test'), 'rb') as cifar100:
            self.data = pickle.load(cifar100, encoding='bytes')
        self.transform = transform

    def __len__(self):
        return len(self.data['data'.encode()])

    def __getitem__(self, index):
        label = self.data['fine_labels'.encode()][index]
        r = self.data['data'.encode()][index, :1024].reshape(32, 32)
        g = self.data['data'.encode()][index, 1024:2048].reshape(32, 32)
        b = self.data['data'.encode()][index, 2048:].reshape(32, 32)
        image = numpy.dstack((r, g, b))

        if self.transform:
            image = self.transform(image)

        return label, image

'''

'''

if train:
    with open(root) as f:
        while True:
            line = f.readline()
            if not line:
                break
            train_content.append(line)
        self.train_data = train_content
else:
    with open(root) as f:
        while True:
            line = f.readline()
            if not line;
               break
            test_content.append(line)
        self.test_data = test_content

'''